import torch
import torch.nn as nn
import torchvision.models as models
from torchvision import transforms
from PIL import Image
import os

class MobileNetV3Classifier(nn.Module):
    def __init__(self, pretrained=True):
        super(MobileNetV3Classifier, self).__init__()
        self.base_model = models.mobilenet_v3_small(pretrained=pretrained)
        in_features = self.base_model.classifier[3].in_features
        self.base_model.classifier[3] = nn.Linear(in_features, 2)  # 2分类（fake/real）

    def forward(self, x):
        return self.base_model(x)

def load_model(model_path, device="cuda" if torch.cuda.is_available() else "cpu"):
    """加载训练好的 MobileNetV3 模型"""
    model = MobileNetV3Classifier(pretrained=False).to(device)
    model.load_state_dict(torch.load(model_path, map_location=device))
    model.eval()  # 设置为评估模式
    return model

def preprocess_image(image_path, transform=None):
    """预处理图像（与训练时一致）"""
    if transform is None:
        transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])
    img = Image.open(image_path).convert("RGB")
    img_tensor = transform(img).unsqueeze(0)  # 增加 batch 维度
    return img_tensor

def predict_image(model, image_path, device="cpu", threshold=0.5):
    """
    使用训练好的 MobileNetV3 预测图像类别
    
    参数:
        model: 加载的模型
        image_path: 图像路径
        device: "cuda" 或 "cpu"
        threshold: 分类阈值（默认0.5）
    
    返回:
        dict: {"label": "Real/Fake", "confidence": float, "is_real": bool}
    """
    # 预处理图像
    img_tensor = preprocess_image(image_path).to(device)
    
    # 预测
    with torch.no_grad():
        output = model(img_tensor)
        probabilities = torch.softmax(output, dim=1)  # 转换为概率
        confidence, predicted_class = torch.max(probabilities, 1)
    
    # 判断类别
    label = "Real Face" if predicted_class.item() == 1 else "Fake Face"
    
    return {
        "label": label,
        "confidence": confidence.item(),
        "is_real": predicted_class.item() == 1,
    }

# 使用示例
if __name__ == "__main__":
    # 1. 加载模型（确保路径正确）
    model_path = "best_mobilenetv3.pth"
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = load_model(model_path, device)
    
    # 2. 预测单张图像
    image_path = "010.jpg"  # 替换为你的测试图像路径
    result = predict_image(model, image_path, device)
    
    # 3. 打印结果
    print(f"预测结果: {result['label']}")
    print(f"置信度: {result['confidence']:.4f}")
    print(f"是否真实人脸: {result['is_real']}")
